import json
import logging
from typing import List, Dict
from openai import OpenAI
from config import Config
import re
import time

logger = logging.getLogger(__name__)

class AgentConfigGenerator:
    def __init__(self, config: Config):
        self.config = config
        self.client = OpenAI(
            api_key=config.DEEPSEEK_API_KEY,
            base_url=config.DEEPSEEK_API_BASE
        )

    # -------------------------------
    # 生成提示词
    # -------------------------------
    def generate_agents_prompt(self, questionnaire_description: str, num_agents: int = 8, existing_agents: List[Dict[str, str]] = None) -> str:
        """生成创建智能体的提示词"""
        existing_part = ""
        if existing_agents:
            existing_desc = "\n".join([f"- {a['name']}：{a['background']}" for a in existing_agents])
            existing_part = f"\n以下智能体已生成，请避免重复或相似：\n{existing_desc}\n"

        return f"""根据以下问卷任务要求，生成{num_agents}个不同背景和性格的智能体配置。要求每个智能体都有独特的背景、性格特点，能够从不同视角回答问卷问题。

问卷任务描述：
{questionnaire_description}
{existing_part}

生成要求：
1. 生成{num_agents}个不同的智能体配置
2. 每个智能体包含：name（名称）、background（背景）、personality（性格特点）
3. 背景要具体，包括年龄、职业/身份、相关经历等
4. 性格特点要鲜明且多样化
5. 确保覆盖不同的年龄段、职业背景、生活经历
6. 避免重复和相似的配置

请严格按照以下JSON格式输出，不要包含其他内容：
{{
    "agents": [
        {{
            "name": "智能体名称",
            "background": "详细的背景描述",
            "personality": "性格特点描述"
        }}
    ]
}}

确保生成{num_agents}个不同的智能体配置。"""

    # -------------------------------
    # 多轮生成主逻辑
    # -------------------------------
    def generate_agents_config(self, questionnaire_description: str, num_agents: int = 8, batch_size: int = 40) -> List[Dict[str, str]]:
        """分批生成智能体配置，自动拼接并去重"""
        all_agents: List[Dict[str, str]] = []
        total_rounds = (num_agents + batch_size - 1) // batch_size

        logger.info(f"开始分{total_rounds}轮生成，总计{num_agents}个智能体配置（每轮{batch_size}个）")

        for i in range(total_rounds):
            remaining = num_agents - len(all_agents)
            if remaining <= 0:
                break

            current_batch = min(batch_size, remaining)
            prompt = self.generate_agents_prompt(
                questionnaire_description,
                current_batch,
                existing_agents=all_agents
            )

            logger.info(f"第{i + 1}/{total_rounds}轮：生成{current_batch}个智能体...")

            for attempt in range(self.config.MAX_RETRIES):
                try:
                    response = self.client.chat.completions.create(
                        model="deepseek-chat",
                        messages=[
                            {
                                "role": "system",
                                "content": "你是一个专业的智能体配置生成器。请严格按照要求的JSON格式输出。"
                            },
                            {"role": "user", "content": prompt}
                        ],
                        temperature=0.8,
                        max_tokens=3000,
                        timeout=self.config.REQUEST_TIMEOUT
                    )

                    response_text = response.choices[0].message.content.strip()
                    batch_agents = self._parse_agents_response(response_text, current_batch)

                    if batch_agents:
                        logger.info(f"✅ 第{i + 1}轮成功生成 {len(batch_agents)} 个智能体")
                        all_agents.extend(batch_agents)
                        break
                    else:
                        logger.warning(f"第{i + 1}轮解析失败，重试中...")
                except Exception as e:
                    logger.warning(f"第{i + 1}轮第{attempt + 1}次失败: {str(e)}")
                    if attempt < self.config.MAX_RETRIES - 1:
                        time.sleep(2)

        # 去重与补全
        all_agents = self._remove_duplicate_agents(all_agents)
        if len(all_agents) < num_agents:
            logger.warning(f"生成数量不足（{len(all_agents)}/{num_agents}），补充默认配置")
            all_agents.extend(self._get_fallback_agents(num_agents - len(all_agents)))

        logger.info(f"最终生成智能体数量：{len(all_agents)}")
        return all_agents[:num_agents]

    # -------------------------------
    # JSON 解析
    # -------------------------------
    def _parse_agents_response(self, response_text: str, expected_count: int) -> List[Dict[str, str]]:
        try:
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                data = json.loads(json_str)

                if 'agents' in data and isinstance(data['agents'], list):
                    agents = data['agents']
                    valid_agents = []
                    for agent in agents:
                        if isinstance(agent, dict) and all(k in agent for k in ['name', 'background', 'personality']):
                            valid_agents.append({
                                'name': str(agent['name']).strip(),
                                'background': str(agent['background']).strip(),
                                'personality': str(agent['personality']).strip()
                            })
                    return valid_agents[:expected_count]

            logger.error("无法从响应中解析出有效的智能体配置")
            return []
        except Exception as e:
            logger.error(f"解析响应失败: {str(e)}")
            return []

    # -------------------------------
    # 去重逻辑
    # -------------------------------
    def _remove_duplicate_agents(self, agents: List[Dict[str, str]]) -> List[Dict[str, str]]:
        seen = set()
        unique_agents = []

        for agent in agents:
            identifier = (
                f"{agent['name']}_{agent['background'][:40]}_{agent['personality'][:20]}"
            )
            # 简单语义相似度判定（模糊匹配）
            simplified_id = re.sub(r'\W+', '', identifier.lower())

            if simplified_id not in seen:
                seen.add(simplified_id)
                unique_agents.append(agent)
            else:
                logger.debug(f"去除重复智能体：{agent['name']}")
        return unique_agents

    # -------------------------------
    # 备用配置
    # -------------------------------
    def _get_fallback_agents(self, num_agents: int) -> List[Dict[str, str]]:
        fallback_agents = [
            {"name": "大学生-理工科", "background": "22岁计算机专业大学生，理性思维强，喜欢数据分析", "personality": "严谨、逻辑性强、注重细节"},
            {"name": "职场白领-市场部", "background": "30岁市场专员，有5年工作经验，擅长沟通", "personality": "外向、有创意、善于表达"},
            {"name": "自由职业者-设计师", "background": "28岁平面设计师，工作灵活，追求生活品质", "personality": "感性、有艺术细胞、独立思考"},
            {"name": "中年管理者", "background": "45岁部门经理，管理经验丰富，注重效率", "personality": "稳重、务实、结果导向"},
            {"name": "退休教师", "background": "65岁退休语文教师，教学经验丰富", "personality": "耐心、细致、喜欢分享"},
            {"name": "创业者", "background": "35岁科技公司创始人，有冒险精神", "personality": "积极、果断、有远见"},
            {"name": "医护人员", "background": "32岁护士，工作认真负责", "personality": "细心、有同理心、专业"},
            {"name": "公务员", "background": "38岁政府工作人员，做事规范", "personality": "谨慎、守规矩、服务意识强"}
        ]
        return fallback_agents[:num_agents]

    # -------------------------------
    # 问卷分析接口（保持原逻辑）
    # -------------------------------
    def analyze_questionnaire_and_generate_agents(self, questions_file: str, num_agents: int = 8) -> List[Dict[str, str]]:
        try:
            with open(questions_file, 'r', encoding='utf-8') as f:
                questions_content = f.read()

            theme_analysis = self._analyze_questionnaire_theme(questions_content)
            questionnaire_description = f"""
问卷内容：
{questions_content}

分析的主题特点：
{theme_analysis}

请根据以上问卷内容和主题分析，生成适合回答这些问题的多样化智能体配置。
"""
            return self.generate_agents_config(questionnaire_description, num_agents, batch_size= 40)
        except Exception as e:
            logger.error(f"分析问卷并生成智能体失败: {str(e)}")
            return self._get_fallback_agents(num_agents)

    def _analyze_questionnaire_theme(self, questions_content: str) -> str:
        analysis_prompt = f"""请分析以下问卷的主要主题和特点，指出需要什么样的受访者背景来有效回答这些问题：

问卷问题：
{questions_content}

请简要总结问卷的主要主题和所需的受访者多样性。"""
        try:
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[{"role": "user", "content": analysis_prompt}],
                temperature=0.3,
                max_tokens=500,
                timeout=self.config.REQUEST_TIMEOUT
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            logger.warning(f"问卷主题分析失败: {str(e)}")
            return "通用社会调查问卷"
